"""
Agent 前端客户端（简化版） - 直接连接到后端API服务

使用同步请求模式，与后端保持一致，
简化了状态管理和通信流程，提供优雅的命令行体验
"""
import requests
import json
import traceback
from typing import Dict, Any, Optional
import time
import uuid
from rich.console import Console
from rich.prompt import Prompt
from rich.panel import Panel
from rich.markdown import Markdown
from rich.theme import Theme
from rich.progress import Progress

# 创建自定义主题
custom_theme = Theme({
    "info": "cyan bold",
    "warning": "yellow bold",
    "success": "green bold",
    "error": "red bold",
    "heading": "magenta bold underline",
    "highlight": "blue bold",
})

# 初始化Rich控制台
console = Console(theme=custom_theme)

# 后端API地址
API_BASE_URL = "http://localhost:8001"

def invoke_agent(user_id: str, query: str, system_message: str = "你会使用工具来帮助用户。如果工具使用被拒绝，请提示用户。"):
    """调用智能体处理查询，并等待完成或中断"""
    # 发送请求到后端API
    payload = {
        "user_id": user_id,
        "query": query,
        "system_message": system_message
    }
    
    console.print("[info]正在发送请求到智能体，请稍候...[/info]")
    
    with Progress() as progress:
        task = progress.add_task("[cyan]处理中...", total=None)
        response = requests.post(f"{API_BASE_URL}/agent/invoke", json=payload)
        progress.update(task, completed=100)
    
    if response.status_code == 200:
        return response.json()
    else:
        raise Exception(f"API调用失败: {response.status_code} - {response.text}")

def resume_agent(user_id: str, session_id: str, response_type: str, args: Optional[Dict[str, Any]] = None):
    """发送响应以恢复智能体执行"""
    payload = {
        "user_id": user_id,
        "session_id": session_id,
        "response_type": response_type,
        "args": args
    }
    
    console.print("[info]正在恢复智能体执行，请稍候...[/info]")
    
    with Progress() as progress:
        task = progress.add_task("[cyan]恢复执行中...", total=None)
        response = requests.post(f"{API_BASE_URL}/agent/resume", json=payload)
        progress.update(task, completed=100)
    
    if response.status_code == 200:
        return response.json()
    else:
        raise Exception(f"恢复智能体执行失败: {response.status_code} - {response.text}")

def get_agent_status(user_id: str):
    """获取智能体状态"""
    response = requests.get(f"{API_BASE_URL}/agent/status/{user_id}")
    
    if response.status_code == 200:
        return response.json()
    else:
        raise Exception(f"获取智能体状态失败: {response.status_code} - {response.text}")

def get_system_info():
    """获取系统信息"""
    response = requests.get(f"{API_BASE_URL}/system/info")
    
    if response.status_code == 200:
        return response.json()
    else:
        raise Exception(f"获取系统信息失败: {response.status_code} - {response.text}")

def delete_agent_session(user_id: str):
    """删除用户会话"""
    response = requests.delete(f"{API_BASE_URL}/agent/session/{user_id}")
    
    if response.status_code == 200:
        return response.json()
    else:
        raise Exception(f"删除会话失败: {response.status_code} - {response.text}")

def process_agent_response(response, user_id):
    """处理智能体响应，包括处理中断和显示结果"""
    session_id = response["session_id"]
    status = response["status"]
    
    # 显示当前会话ID
    console.print(f"[info]当前会话ID: {session_id}[/info]")
    
    # 处理不同状态
    if status == "interrupted":
        # 处理中断
        interrupt_data = response.get("interrupt_data", {})
        
        # 显示中断信息
        message = interrupt_data.get("message", "需要您的输入")
        console.print(Panel(
            f"{message}",
            title="[warning]智能体需要您的决定[/warning]",
            border_style="yellow"
        ))
        
        # 获取用户输入
        user_input = Prompt.ask("[highlight]您的选择[/highlight]")
        
        # 处理用户输入
        try:
            while True:
                if user_input.lower() == "yes":
                    response = resume_agent(user_id, session_id, "accept")
                    break
                elif user_input.lower() == "no":
                    response = resume_agent(user_id, session_id, "reject")
                    break
                elif user_input.lower() == "edit":
                    # 获取新的查询内容
                    new_query = Prompt.ask("[highlight]请输入新的搜索内容[/highlight]")
                    response = resume_agent(user_id, session_id, "edit", args={"query": new_query})
                    break
                else:
                    console.print("[error]无效输入，请输入 'yes'、'no' 或 'edit'[/error]")
                    user_input = Prompt.ask("[highlight]您的选择[/highlight]")
            
            # 重新获取用户输入（维持当前响应不变）
            return process_agent_response(response, user_id)
        
        except Exception as e:
            console.print(f"[error]处理响应时出错: {str(e)}[/error]")
            return None
    
    elif status == "completed":
        # 显示结果
        result = response.get("result", {})
        if result and "messages" in result:
            final_message = result["messages"][-1]
            console.print(Panel(
                Markdown(final_message["content"]),
                title="[success]智能体回答[/success]",
                border_style="green"
            ))
        else:
            console.print("[warning]智能体没有返回有效的消息[/warning]")
        
        # 会话完成后自动删除会话
        try:
            delete_response = delete_agent_session(user_id)
            console.print(f"[success]会话已结束并清理 (Session ID: {session_id})[/success]")
        except Exception as e:
            console.print(f"[warning]删除会话时出错: {str(e)}[/warning]")
        
        return result
    
    elif status == "error":
        # 显示错误信息
        error_msg = response.get("message", "未知错误")
        console.print(Panel(
            f"{error_msg}",
            title="[error]处理过程中出错[/error]",
            border_style="red"
        ))
        
        # 错误状态下也删除会话
        try:
            delete_response = delete_agent_session(user_id)
            console.print(f"[info]会话已清理 (Session ID: {session_id})[/info]")
        except Exception as e:
            console.print(f"[warning]删除会话时出错: {str(e)}[/warning]")
        
        return None
    
    else:
        # 其他未知状态
        console.print(f"[error]智能体状态: {status} - {response.get('message', '无消息')}[/error]")
        return response

def main():
    """主函数，运行客户端"""
    console.print(Panel(
        "Agent 前端客户端（简化版）",
        title="[heading]智能体交互系统[/heading]",
        border_style="magenta"
    ))
    
    # 尝试获取系统信息
    try:
        system_info = get_system_info()
        console.print(f"[info]系统活跃会话: {system_info['sessions_count']}[/info]")
    except Exception:
        console.print("[warning]无法获取系统状态，但这不影响使用[/warning]")
    
    # 获取用户ID（在真实应用中可能是登录后的用户标识）
    default_user_id = f"user_{int(time.time())}"
    user_id = Prompt.ask("[info]请输入用户ID[/info]", default=default_user_id)
    
    while True:
        try:
            # 获取用户查询
            query = Prompt.ask("\n[info]请输入您的问题[/info] (输入 'exit' 退出，'status' 查看会话状态)", 
                               default="默认：你好！")
            
            if query.lower() == 'exit':
                console.print("[info]感谢使用，再见！[/info]")
                break
            elif query.lower() == 'status':
                # 查看会话状态
                try:
                    status_info = get_agent_status(user_id)
                    console.print(Panel(
                        f"用户ID: {status_info.get('user_id', 'N/A')}\n"
                        f"会话ID: {status_info.get('session_id', 'N/A')}\n"
                        f"状态: {status_info.get('status', 'N/A')}",
                        title="[info]当前会话状态[/info]",
                        border_style="cyan"
                    ))
                except Exception as e:
                    console.print(f"[warning]获取会话状态失败: {str(e)}[/warning]")
                continue
                
            # 调用智能体
            response = invoke_agent(user_id, query)
            
            # 处理智能体响应
            process_agent_response(response, user_id)
            
        except KeyboardInterrupt:
            console.print("\n[warning]用户中断，正在退出...[/warning]")
            break
        except Exception as e:
            console.print(f"[error]运行过程中出错: {str(e)}[/error]")
            console.print(traceback.format_exc())
            continue

if __name__ == "__main__":
    main()
